# -*- coding: utf-8 -*-
"""
    codegen
    ~~~~~~~

    Extension to ast that allow ast -> CPP code generation.

    :copyright: Copyright 2008 by Armin Ronacher (adapted for C++ by Jurgis Pralgauskis)
    :license: BSD.

    https://pypi.python.org/pypi/codegen

    also similar: http://svn.python.org/view/python/trunk/Demo/parser/unparse.py?view=markup
    and "the missing Python AST docs" http://greentreesnakes.readthedocs.org/en/latest/ 
"""
from ast import *


BOOLOP_SYMBOLS = {
    And:        'and', # &&
    Or:         'or'   # ||
}

BINOP_SYMBOLS = {
    Add:        '+',
    Sub:        '-',
    Mult:       '*',
    Div:        '/',
    FloorDiv:   '//',
    Mod:        '%',
    LShift:     '<<',
    RShift:     '>>',
    BitOr:      '|',
    BitAnd:     '&',
    BitXor:     '^',
    #~ Pow:     '**'  # in cpp should translate as function
}

CMPOP_SYMBOLS = {
    Eq:         '==',
    Gt:         '>',
    GtE:        '>=',
    Lt:         '<',
    LtE:        '<=',
    NotEq:      '!=',
    In:         'IN /*???*//*!!!*/',
    #~ Is:         'is',
    #~ IsNot:      'is not',
    #~ NotIn:      'not in'
}

UNARYOP_SYMBOLS = {
    #~ Invert:     '~',
    Not:        '!',
    UAdd:       '+',
    USub:       '-'
}

ALL_SYMBOLS = {}
ALL_SYMBOLS.update(BOOLOP_SYMBOLS)
ALL_SYMBOLS.update(BINOP_SYMBOLS)
ALL_SYMBOLS.update(CMPOP_SYMBOLS)
ALL_SYMBOLS.update(UNARYOP_SYMBOLS)


def to_source(node, indent_with=' ' * 4, add_line_information=False, comments={}):
    """This function can convert a node tree back into python sourcecode.
    This is useful for debugging purposes, especially if you're dealing with
    custom asts not generated by python itself.

    It could be that the sourcecode is evaluable when the AST itself is not
    compilable / evaluable.  The reason for this is that the AST contains some
    more data than regular sourcecode does, which is dropped during
    conversion.

    Each level of indentation is replaced with `indent_with`.  Per default this
    parameter is equal to four spaces as suggested by PEP 8, but it might be
    adjusted to match the application's styleguide.

    If `add_line_information` is set to `True` comments for the line numbers
    of the nodes are added to the output.  This can be used to spot wrong line
    number information of statement nodes.
    """
    generator = SourceGenerator(indent_with, add_line_information, comments)
    generator.visit(node)
    #~ print (generator.result)
    return ''.join(generator.result)


class SourceGenerator(NodeVisitor):
    """This visitor is able to transform a well formed syntax tree into python
    sourcecode.  For more details have a look at the docstring of the
    `node_to_source` function.
    """

    def __init__(self, indent_with=' '*4, add_line_information=False, comments_dict={}):
        self.result = []
        self.indent_with = indent_with
        self.add_line_information = add_line_information
        self.indentation = 0
        self.new_lines = 0
        self.defined_variables = {} # TODO: scopes
        #~ self.variable_values = {} # TODO: scopes
        self.varaible2indexed_substitutions = {} # TODO: scopes
        self.opstack=[]
        self.comments=comments_dict # dictionary
        self.commentPending =  None
        

    def dump_node(self, node):
        self.write('\n')
        n = len(self.result)
        #~ newline_end = ['\n'] if self.result[-1] == '\n' else []
            
        self.visit(node)
        result = self.result[n:]
        self.result = self.result[:n]
        return ''.join(result).lstrip('\n')
    
    def write(self, x):
        if self.new_lines:
            # inject comments -- before next newline
            if self.commentPending and self.commentLineEnd:
                self.result.append(' //'+self.commentPending )
                self.commentPending =  None
            self.commentLineEnd = True  
            
            if self.result:
                self.result.append('\n' * self.new_lines)
            self.result.append(self.indent_with * self.indentation)
            self.new_lines = 0
        self.result.append(x)

    def newline(self, node=None, extra=0):
        self.new_lines = max(self.new_lines, 1 + extra)
        if hasattr(node, 'lineno') and node.lineno in self.comments:
            self.commentPending = self.comments[node.lineno]  # we "pass" it via commentPending, because in "write" we don't have "lineno"
            self.commentLineEnd = False
            #~ self.write('// ' +str(node.lineno) +'  ' + self.comments[node.lineno])
            #~ self.new_lines = 1
        if node is not None and self.add_line_information:
            self.write('// line: %s' % node.lineno)
            self.new_lines = 1


    def body(self, statements):
        self.new_line = True
        self.indentation += 1
        for stmt in statements:
            self.visit(stmt)
        self.indentation -= 1

    def body_or_else(self, node):
        self.body(node.body)
        if node.orelse:
            self.newline()
            self.write('else:')
            self.body(node.orelse)

    def signature(self, node):
        want_comma = []
        def write_comma():
            if want_comma:
                self.write(', ')
            else:
                want_comma.append(True)

        padding = [None] * (len(node.args) - len(node.defaults))
        for arg, default in zip(node.args, padding + node.defaults):
            write_comma()
            
            self.write_type() # int
            #~ self.visit(arg)
            self.write(arg.arg) # changed in py3: arg.id is empty
            if default is not None:
                self.write('=')
                self.visit(default)
        if node.vararg is not None:
            write_comma()
            self.write('*' + node.vararg)
        if node.kwarg is not None:
            write_comma()
            self.write('**' + node.kwarg)

    def decorators(self, node):
        for decorator in node.decorator_list:
            self.newline(decorator)
            self.write('@')
            self.visit(decorator)

    # Statements

    #~ def visit_Arg(self, node):
        #~ # arg = (identifier arg, expr? annotation)
        #~ self.write(node.arg)
        
    def visit_Assign(self, node):
        self.newline(node)
        for idx, target in enumerate(node.targets):
            if idx:
                self.write(', ')
            #~ self.visit(target)
            self.visit_VariableDef_cpp(target, node.value)

        self.write(' = ')
        self.visit(node.value)
        self.write(';')

    def visit_AugAssign(self, node):
        self.newline(node)
        self.visit(node.target)
        self.write(BINOP_SYMBOLS[type(node.op)] + '=')
        self.visit(node.value)
        self.write(';')

    def visit_ImportFrom(self, node):
        self.newline(node)
        self.write('from %s%s import ' % ('.' * node.level, node.module))
        for idx, item in enumerate(node.names):
            if idx:
                self.write(', ')
            self.write(item)

    def visit_Import(self, node):
        self.newline(node)
        for item in node.names:
            self.write('import ')
            self.visit(item)

    def visit_Expr(self, node):
        self.newline(node)
        self.generic_visit(node)

    def visit_FunctionDef(self, node):
        self.newline(extra=1)
        self.decorators(node)
        self.newline(node)
        self.write('int %s(' % node.name)  # be carefull with other <type>s
        self.signature(node.args)
        self.write('){')
        self.body(node.body)
        self.newline(node)
        self.write('}')
        

    def visit_ClassDef(self, node):
        have_args = []
        def paren_or_comma():
            if have_args:
                self.write(', ')
            else:
                have_args.append(True)
                self.write('(')

        self.newline(extra=2)
        self.decorators(node)
        self.newline(node)
        self.write('class %s' % node.name)
        for base in node.bases:
            paren_or_comma()
            self.visit(base)
        # XXX: the if here is used to keep this module compatible
        #      with python 2.6.
        if hasattr(node, 'keywords'):
            for keyword in node.keywords:
                paren_or_comma()
                self.write(keyword.arg + '=')
                self.visit(keyword.value)
            if node.starargs is not None:
                paren_or_comma()
                self.write('*')
                self.visit(node.starargs)
            if node.kwargs is not None:
                paren_or_comma()
                self.write('**')
                self.visit(node.kwargs)
        self.write(have_args and '):' or ':')
        self.body(node.body)

    def visit_If(self, node):
        self.newline(node)
        self.write('if (')
        self.visit(node.test)
        self.write('){')
        self.body(node.body)
        self.newline(node);  self.write('}')  # finish block for cpp
        while True:
            else_ = node.orelse
            if len(else_) == 1 and isinstance(else_[0], If):
                node = else_[0]
                self.newline()
                self.write('else if ')
                self.visit(node.test)
                self.write('{')
                self.body(node.body)
                self.newline(node);  self.write('}')  # finish block for cpp
            else:
                if else_:
                    self.newline()
                    self.write('else {')
                    self.body(else_)
                    self.newline(node);  self.write('}')  # finish block for cpp
                break

    def visit_VariableDef_cpp(self, target, value):
        if isinstance(target, Name):    # variable definition. TODO: arrays indicators...
            if not target.id in self.defined_variables:
                #~ self.defined_variables.append(target.id)
                self.defined_variables[target.id] = value
                self.write_type(value)
                self.visit(target)
                self.write_dimension(value)
            else: self.visit(target)
        else: self.visit(target)
            
    def write_dimension(self, value):
        if isinstance(value, List): 
            self.write('[%s]'%len(value.elts) )
            # for multidimentional - Matrixes
            if len(value.elts) > 0: # TODO # could do recursively:
                value_inner = value.elts[0]
                if isinstance(value_inner, List):
                    max_len = max( [len(x.elts) for x in value.elts if isinstance(x, List)] )
                    self.write('[%s]' % max_len)
                
            
    def write_type(self, value=0):
        if isinstance(value, Num):
            self.write('int ')
        elif isinstance(value, Str):
            self.write('string ')            
        elif isinstance(value, List):
            #~ self.write('int []')
            if len(value.elts) > 0: # could do recursively:
                value = value.elts[0]
                if isinstance(value, Num):
                    self.write('int ')
                elif isinstance(value, Str):
                    self.write('string ')
                else: self.write('int ') # default
            else: self.write('int ') # default
        else: self.write('int ') # default

    
    
    def visit_For(self, node):
        self.newline(node)
        
        if isinstance (node.iter, Call) and node.iter.func.id == 'range':
            py_range_args = [ x.n if isinstance(x, Num)
                              else self.dump_node(x)
                            for x in node.iter.args
                           ]
            if len(py_range_args) == 1:
                cpp_range_args = [0, py_range_args[0], 1]
            if len(py_range_args) == 2:
                cpp_range_args = [py_range_args[0], py_range_args[1], 1]
            if len(py_range_args) == 3:
                cpp_range_args = py_range_args

            self.write('for (')
            self.write_type() # default int 
            self.visit(node.target)
            self.write('=%s; ' % cpp_range_args[0])
            self.visit(node.target)
            
            if cpp_range_args[2] > 0:  # if ascending
                self.write( '<%s; ' % cpp_range_args[1])
                self.visit(node.target)
                if cpp_range_args[2] == 1:
                    self.write('++ ')
                else:
                    self.write('+=%s ' % cpp_range_args[2])
                    
            else:                   # if descending 
                self.write( ' >%s; ' % cpp_range_args[1])
                self.visit(node.target)
                if cpp_range_args[2] == -1:
                    self.write('-- ')
                else:
                    self.write('-=%s ' % -cpp_range_args[2])
                    
            self.write('){')
        else:
            thelist = node.iter
            # raw list
            if isinstance (thelist, List):
                # find (unused) name for list/array
                for arrname in "ABCDEFGHIJKLMNPRSTVWXZ":
                    if not arrname in self.defined_variables:
                        self.defined_variables[arrname] = thelist
                        break
                        
                self.write_type(thelist)
                self.write(arrname)
                self.write_dimension(thelist)
                self.write( ' = %s; ' % thelist.elts )
                self.newline()
                
                looplen = len( thelist.elts ) 
                
            # list name -- if we find name, then it should have been decleared earlier
            if isinstance (thelist, Name):
                arrname = thelist.id
                arrvalues = self.defined_variables.get(arrname, None )
                looplen = '%s /*???*/' % len( arrvalues.elts if arrvalues else [])   # ??? means warning -- as not sure if the current value is the last assigned to variable
                
            for loopvar in "ijkabcdxyzw":
                if not loopvar in self.defined_variables:
                    self.defined_variables[loopvar] = None
                    break
            self.varaible2indexed_substitutions[node.target] = "%s[%s]" % (arrname, loopvar);
            
            
            self.write('for (int @i=0; @i<@len; @i++){'.replace("@i", loopvar).replace('@len', str(looplen) ) )
            self.newline()
            #~ self.write('    %s = %s;' % (node.target.id, self.varaible2indexed_substitutions[node.target]) ) 
            self.write('    ')
            self.write('int ')  # if node.target is simple variable -- most cases -- it should be declared
            self.visit(node.target) # might be expression or so   # if it is -- what to do with definition?
            self.write(' = %s;' % ( self.varaible2indexed_substitutions[node.target]) ) 

            
        self.body_or_else(node)
        self.newline(node);  self.write('}')  # finish block for cpp

    def visit_While(self, node):
        self.newline(node)
        self.write('while (')
        self.visit(node.test)
        self.write('){')
        self.body_or_else(node)
        self.newline(node);  self.write('}')        

    def visit_With(self, node):
        self.newline(node)
        self.write('with ')
        self.visit(node.context_expr)
        if node.optional_vars is not None:
            self.write(' as ')
            self.visit(node.optional_vars)
        self.write(': // !!!!!!!!!!!!!!!!!!!! ')  # should manually change in other langs
        self.body(node.body)
        

    def visit_Pass(self, node):
        self.newline(node)
        self.write('pass')

    def visit_Print(self, node):
        # XXX: python 2.6 only
        #~ self.newline(node)        
        self.visit_Print_call(node.values, node=node)

    def visit_Print_call(self, values, node=None, nl=True):

        self.newline(node)
        self.write('cout << ')
        want_comma = False
        for value in values:
            if want_comma:
                self.write(' << " " << ')
            self.visit(value)
            want_comma = True
        if nl:
            self.write(' << endl')
        self.write(';')

    def visit_Delete(self, node):
        self.newline(node)
        #~ self.write('del ')
        #TypeError: 'Delete' object is not iterable
        for idx, target in enumerate(node.targets):
            if idx:
                self.write(', ')
            self.visit(target)
        self.write(' = 0; /*???*/')

    def visit_TryExcept(self, node):
        self.newline(node)
        self.write('try:')
        self.body(node.body)
        for handler in node.handlers:
            self.visit(handler)

    def visit_TryFinally(self, node):
        self.newline(node)
        self.write('try:')
        self.body(node.body)
        self.newline(node)
        self.write('finally:')
        self.body(node.finalbody)

    def visit_Global(self, node):
        self.newline(node)
        self.write('global ' + ', '.join(node.names))

    def visit_Nonlocal(self, node):
        self.newline(node)
        self.write('nonlocal ' + ', '.join(node.names))

    def visit_Return(self, node):
        self.newline(node)
        self.write('return ')
        self.visit(node.value)
        self.write(';')

    def visit_Break(self, node):
        self.newline(node)
        self.write('break ;')

    def visit_Continue(self, node):
        self.newline(node)
        self.write('continue ;')

    def visit_Raise(self, node):
        # XXX: Python 2.6 / 3.0 compatibility
        self.newline(node)
        self.write('raise')
        if hasattr(node, 'exc') and node.exc is not None:
            self.write(' ')
            self.visit(node.exc)
            if node.cause is not None:
                self.write(' from ')
                self.visit(node.cause)
        elif hasattr(node, 'type') and node.type is not None:
            self.visit(node.type)
            if node.inst is not None:
                self.write(', ')
                self.visit(node.inst)
            if node.tback is not None:
                self.write(', ')
                self.visit(node.tback)

    # Expressions

    def visit_Attribute(self, node):
        self.visit(node.value)
        self.write('.' + node.attr)

    def visit_Call(self, node):
        want_comma = []
        def write_comma():
            if want_comma:
                self.write(', ')
            else:
                want_comma.append(True)

        process_normally = False
        if isinstance(node.func, Name) and node.func.id == 'print':
            self.visit_Print_call(values=node.args, node=node.func)
        elif isinstance(node.func, Name) and node.func.id == 'len':
            if isinstance(node.args[0], List):
                length = len( node.args[0].elts )
                self.write( length )
            elif isinstance(node.args[0], Name):
                arrvalues = self.defined_variables.get(node.args[0].id, None )
                length = '%s /*???*/' % len( arrvalues.elts if arrvalues else [])   # ??? means warning -- as not sure if the current value is the last assigned to variable
                self.write( length )
            else:
                self.write( '/*???*/' )
                process_normally = True
        else:
            process_normally = True

        if process_normally:
            self.visit(node.func)
            self.write('(')
            for arg in node.args:
                write_comma()
                self.visit(arg)
            for keyword in node.keywords:
                write_comma()
                self.write(keyword.arg + '=')
                self.visit(keyword.value)
            if node.starargs is not None:
                write_comma()
                self.write('*')
                self.visit(node.starargs)
            if node.kwargs is not None:
                write_comma()
                self.write('**')
                self.visit(node.kwargs)
            self.write(')')

    def visit_Name(self, node):
        self.write(node.id)

    def visit_Str(self, node):
        #~ if len(node.s) == 1: # char ??
             #~ self.write(repr(node.s))
        self.write(repr(node.s).replace("'", '"'))  # double quote for string
        

    def visit_Bytes(self, node):
        self.write(repr(node.s))

    def visit_Num(self, node):
        self.write(repr(node.n))

    def visit_Tuple(self, node):
        self.write('(')
        idx = -1
        for idx, item in enumerate(node.elts):
            if idx:
                self.write(', ')
            self.visit(item)
        self.write(idx and ')' or ',)')

    def sequence_visit(left, right):
        def visit(self, node):
            multidimentional = False
            #~ multidimentional_indent = len( self.result[-1] ) + 2
            self.write(left)
            
            for idx, item in enumerate(node.elts):
                if idx:
                    self.write(', ')
                
                if isinstance(item, List):  # for matrix
                    multidimentional = True
                    self.newline(item); self.write(' '*17)
                                   
                self.visit(item)
            if multidimentional:
                self.newline(node); self.write(' '*14)  
            self.write(right)
        return visit

    visit_List = sequence_visit('{', '}')
    #~ visit_Set = sequence_visit('{', '}')
    del sequence_visit

    def visit_Dict(self, node):
        self.write('{')
        for idx, (key, value) in enumerate(zip(node.keys, node.values)):
            if idx:
                self.write(', ')
            self.visit(key)
            self.write(': ')
            self.visit(value)
        self.write('}')


    def dump_opstack(self):
        #~ print(self.opstack)
        pass
        
    def visit_BinOp(self, node):
        
        if type(node.op) == Pow:
            self.write('pow(')
            self.visit(node.left)
            self.write(',')
            self.visit(node.right)
            self.write(')')
        else:
            if type(node.op) in [Add, Sub] and self.opstack:   self.write('(') # TODO: decide when they are needed, when -- not...
            self.opstack.append(str(node.op));  self.dump_opstack()
            self.visit(node.left)
            #~ if type(node.op) == Div: # might need 1.0* in front to become double/float
            if type(node.op) == FloorDiv:
                self.write('/') # cpp doesn't have //
            else:
                self.write(' %s ' % BINOP_SYMBOLS[type(node.op)])
            self.visit(node.right)
            del self.opstack[-1]
            if type(node.op) in [Add, Sub] and self.opstack:   self.write(')') # TODO: decide when they are needed, when -- not...

    def visit_BoolOp(self, node):
        if type(node.op) == Or and self.opstack:  self.write('(')
        self.opstack.append(str(node.op));  self.dump_opstack()
        for idx, value in enumerate(node.values):
            if idx:
                self.write(' %s ' % BOOLOP_SYMBOLS[type(node.op)])
            self.visit(value)
        del self.opstack[-1]
        if type(node.op) == Or and self.opstack:  self.write(')')

    def visit_Compare(self, node):
        if self.opstack:  self.write('(')
        self.opstack.append('..'.join([str(op) for op in node.ops]));  self.dump_opstack()
        self.visit(node.left)
        for op, right in zip(node.ops, node.comparators):
            self.write(' %s ' % CMPOP_SYMBOLS[type(op)])
            self.visit(right)
        del self.opstack[-1]
        if self.opstack: self.write(')')

    def visit_UnaryOp(self, node):
        if self.opstack: self.write('(')
        self.opstack.append(str(node.op));  self.dump_opstack()
        op = UNARYOP_SYMBOLS[type(node.op)]
        self.write(op)
        if op == 'not':
            self.write(' ')
        self.visit(node.operand)
        del self.opstack[-1]
        if self.opstack: self.write(')')

    def visit_Subscript(self, node):
        self.visit(node.value)
        self.write('[')
        self.visit(node.slice)
        self.write(']')

    def visit_Slice(self, node):
        if node.lower is not None:
            self.visit(node.lower)
        self.write(':')
        if node.upper is not None:
            self.visit(node.upper)
        if node.step is not None:
            self.write(':')
            if not (isinstance(node.step, Name) and node.step.id == 'None'):
                self.visit(node.step)

    def visit_ExtSlice(self, node):
        for idx, item in node.dims:
            if idx:
                self.write(', ')
            self.visit(item)

    def visit_Yield(self, node):
        self.write('yield ')
        self.visit(node.value)

    def visit_Lambda(self, node):
        self.write('lambda ')
        self.signature(node.args)
        self.write(': ')
        self.visit(node.body)

    def visit_Ellipsis(self, node):
        self.write('Ellipsis')

    def generator_visit(left, right):
        def visit(self, node):
            self.write(left)
            self.visit(node.elt)
            for comprehension in node.generators:
                self.visit(comprehension)
            self.write(right)
        return visit

    visit_ListComp = generator_visit('[', ']')
    visit_GeneratorExp = generator_visit('(', ')')
    visit_SetComp = generator_visit('{', '}')
    del generator_visit

    def visit_DictComp(self, node):
        self.write('{')
        self.visit(node.key)
        self.write(': ')
        self.visit(node.value)
        for comprehension in node.generators:
            self.visit(comprehension)
        self.write('}')

    def visit_IfExp(self, node):
        self.visit(node.body)
        self.write(' if ')
        self.visit(node.test)
        self.write(' else ')
        self.visit(node.orelse)

    def visit_Starred(self, node):
        self.write('*')
        self.visit(node.value)

    def visit_Repr(self, node):
        # XXX: python 2.6 only
        self.write('`')
        self.visit(node.value)
        self.write('`')

    # Helper Nodes

    def visit_alias(self, node):
        self.write(node.name)
        if node.asname is not None:
            self.write(' as ' + node.asname)

    def visit_comprehension(self, node):
        self.write(' for ')
        self.visit(node.target)
        self.write(' in ')
        self.visit(node.iter)
        if node.ifs:
            for if_ in node.ifs:
                self.write(' if ')
                self.visit(if_)

    def visit_excepthandler(self, node):
        self.newline(node)
        self.write('except')
        if node.type is not None:
            self.write(' ')
            self.visit(node.type)
            if node.name is not None:
                self.write(' as ')
                self.visit(node.name)
        self.write(':')
        self.body(node.body)
